/*
 * Copyright (c) 2025 [ByteDance Ltd. and/or its affiliates.] and/or its affiliates
 * SPDX-License-Identifier: LGPL v2.1+
*/

#include "FFmpegBaseIOCryptor.h"

extern "C" {
#include "libavformat/avio.h"
}

FFmpegBaseIOCryptor::~FFmpegBaseIOCryptor() {
    free();
}

void FFmpegBaseIOCryptor::free() {
    if (m_pOriginalIOCtx && m_bOwnOriginalIOCtx) {
        avio_close(m_pOriginalIOCtx);
        if (m_pFormatContext) {
            m_pFormatContext->pb = nullptr;
        }
    } else {
        if (m_pFormatContext) {
            m_pFormatContext->pb = m_pOriginalIOCtx;
        }
    }

    if (m_pCryptIOCtx) {
        if (m_pCryptIOCtx->buffer) {
            av_freep(&m_pCryptIOCtx->buffer);
        }
        m_pCryptIOCtx->buffer = nullptr;
        av_free(m_pCryptIOCtx);
    }

    m_bOwnOriginalIOCtx = false;
    m_pFormatContext = nullptr;
    m_pOriginalIOCtx = nullptr;
    m_pCryptIOCtx = nullptr;
}

int FFmpegBaseIOCryptor::prepare(const char *url, AVFormatContext *formatContext, const AVDictionary *options, const std::map<std::string, FFAny> *params) {
    free();

    if (params) {
        auto ret = parseParam(params);
        if (ret < 0)
            return ret;
    }

    if (formatContext->pb) {
        m_pOriginalIOCtx = formatContext->pb;
        m_bOwnOriginalIOCtx = false;
    } else {
        AVDictionary *tmpOptions = nullptr;
        av_dict_copy(&tmpOptions, options, 0);
        AVIOContext *context = nullptr;
        auto ret = avio_open2(&context, url, isEncryptor() ? AVIO_FLAG_READ_WRITE : AVIO_FLAG_READ, nullptr, &tmpOptions);
        av_dict_free(&tmpOptions);
        if (ret < 0)
            return ret;

        m_pOriginalIOCtx = context;
        m_bOwnOriginalIOCtx = true;
    }

    auto buffer = (uint8_t *)av_malloc(m_pOriginalIOCtx->buffer_size);
    m_pCryptIOCtx = avio_alloc_context(buffer, m_pOriginalIOCtx->buffer_size, 
                                       isEncryptor() ? 1 : 0, 
                                       this, 
                                       FFmpegBaseIOCryptor::ReadPacket,
                                       FFmpegBaseIOCryptor::WritePacket,
                                       FFmpegBaseIOCryptor::Seek);
    formatContext->pb = m_pCryptIOCtx;
    formatContext->flags |= AVFMT_FLAG_CUSTOM_IO;
    m_pFormatContext = formatContext;
    return 0;
}

void FFmpegBaseIOCryptor::clearFormatContext() {
    m_pFormatContext = nullptr;
}

int FFmpegBaseIOCryptor::ReadPacket(void *opaque, uint8_t *buf, int bufSize) {
    auto context = static_cast<FFmpegBaseIOCryptor *>(opaque);
    return context->readPacket(buf, bufSize);
}

int FFmpegBaseIOCryptor::WritePacket(void *opaque, const uint8_t *buf, int bufSize) {
    auto context = static_cast<FFmpegBaseIOCryptor *>(opaque);
    return context->writePacket(buf, bufSize);
}

int64_t FFmpegBaseIOCryptor::Seek(void *opaque, int64_t offset, int whence) {
    auto context = static_cast<FFmpegBaseIOCryptor *>(opaque);
    return context->seek(offset, whence);
}

int64_t FFmpegBaseIOCryptor::seek(int64_t offset, int whence) {
    return m_pOriginalIOCtx->seek(m_pOriginalIOCtx->opaque, offset, whence);
}
